package com.chiancloud.shield.core.main;

import org.apache.hadoop.service.Service;
import org.apache.hadoop.util.ShutdownHookManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.ref.WeakReference;

public class ServiceShutdownHook implements Runnable {
  private static final Logger LOG = LoggerFactory.getLogger(
      ServiceShutdownHook.class);

  private final WeakReference<Service> serviceRef;
  private Runnable hook;

  public ServiceShutdownHook(Service service) {
    serviceRef = new WeakReference<Service>(service);
  }

  public void register(int priority) {
    unregister();
    hook = this;
    ShutdownHookManager.get().addShutdownHook(hook, priority);
  }

  public synchronized void unregister() {
    if (hook != null) {
      try {
        ShutdownHookManager.get().removeShutdownHook(hook);
      } catch (IllegalStateException e) {
        LOG.info("Failed to unregister shutdown hook: {}", e, e);
      }
      hook = null;
    }
  }

  @Override
  public void run() {
    Service service;
    synchronized (this) {
      service = serviceRef.get();
      serviceRef.clear();
    }
    if (service == null) {
      return;
    }
    try {
      service.stop();
    } catch (Throwable t) {
      LOG.info("Error stopping {}: {}", service.getName(), t);
    }
  }
}
